# -*- coding:utf-8 -*-
import asyncio
from fastapi import FastAPI, Request
from fastapi.middleware.trustedhost import TrustedHostMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware


class CustomHeaderMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        response = await call_next(request)
        response.headers['Custom'] = 'Example'
        return response


async def timeout_middleware(req: Request, call_next, timeout):
    try:
        return await asyncio.wait_for(call_next(req), timeout=timeout)
    except asyncio.TimeoutError:
        return JSONResponse(
            {'code': 500, 'msg': "超时", 'data': []})


def register_middleware(app: FastAPI, settings) -> None:
    """
    注册中间件
    """

    # 跨域
    app.add_middleware(
        CORSMiddleware,
        allow_origins=settings.CORS_ORIGINS,
        allow_credentials=settings.CORS_ALLOW_CREDENTIALS,
        allow_methods=settings.CORS_ALLOW_METHODS,
        allow_headers=settings.CORS_ALLOW_HEADERS,
    )

    # 允许的 hosts 列表
    app.add_middleware(TrustedHostMiddleware, allowed_hosts=settings.AllowHosts)

    # 处理包含”gzip”在Accept-Encoding标头中的任何请求的 GZip 响应。
    app.add_middleware(GZipMiddleware, minimum_size=1000)

    # app.add_middleware(ProxyHeadersMiddleware)

    app.add_middleware(CustomHeaderMiddleware)
    # 自定义中间件- 连接超时
    # @app.middleware("http")
    # async def timeout1(req:Request, call_next):

    #     resp = await timeout_middleware(req, call_next, 5)
    #     print("TimeOut1-middleware")

    #     return resp

    # 自定义中间件- 访问频度

    # 自定义中间件- 敏感词过滤

    return
